
import os
import sys
import cv2
import time
import json
import numpy as np
from deeplabplus.tensorrt_infer import TensorRTInfer

def infer_video():
    with open("deeplabplus/cfg/deeplabv3plus_r18_classout.json", 'r') as file:
        cfg = json.load(file)

    print(f"path = {cfg['engine_pth']}, infer = {cfg['infer']}" )
    
    model_infer = TensorRTInfer(cfg["engine_pth"], cfg["infer"])
    model_infer.initialize()

    video_path = 'test.mp4'
    output_path = 'test_classsout_0703.mp4'
    sample_frame = 2
    save_img = False

    cap = cv2.VideoCapture(video_path)

    output_fourcc = 'mp4v'
    assert (cap.isOpened())
    input_fps = cap.get(cv2.CAP_PROP_FPS)
    # init output video
    fourcc = cv2.VideoWriter_fourcc(*output_fourcc)
    output_fps = input_fps
    output_height = 360
    output_width = 640
    writer = cv2.VideoWriter(
        output_path, fourcc, output_fps, (output_width, output_height), True)

    # start looping
    cnt = -1
    times = []
    pre_ts = []
    post_ts = []
    try:
        while True:
            flag, frame = cap.read()
            if not flag:
                break
            cnt += 1
            if cnt % sample_frame != 0:
                continue

            # test a single image
            t = time.time()
            result = model_infer.infer(frame)
            dur = round((time.time() - t) * 1000, 3)
            print("整体耗时：", dur)
            times.append(dur)

            draw_img_post = model_infer.visualize(frame, result)

            if writer:
                if draw_img_post.shape[0] != output_height or draw_img_post.shape[
                        1] != output_width:
                    draw_img_post = cv2.resize(draw_img_post,
                                          (output_width, output_height))
                writer.write(cv2.cvtColor(draw_img_post, cv2.COLOR_RGB2BGR))
    except KeyboardInterrupt:
        pass

    writer.release()
    cap.release()
    model_infer.clean()

    print("————————\n平均整体耗时：", np.mean(times))
    print("平均前向耗时：", np.mean(pre_ts))
    print("平均后向耗时：", np.mean(post_ts))
    model_infer.clean()

if __name__ == '__main__':
    infer_video()